import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import matplotlib.pyplot as plt
plt.rcParams['font.sans-serif'] = ['SimHei']  # 设置中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决保存图像时负号'-'显示为方块的问题

def getTrainSetAndTestSet(DataPath):
    data = pd.read_csv(DataPath)
    X = data[['AT', 'V', 'AP', 'RH']]
    y = data[['PE']]
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=1)

    print("\n数据集维度：")
    print("原始数据特征维度:", X.shape)
    print("训练集特征维度:", X_train.shape)
    print("测试集特征维度:", X_test.shape)
    print("训练集标签维度:", y_train.shape)
    print("测试集标签维度:", y_test.shape)

    return X_train, X_test, y_train, y_test

def TrainLinearRegression(X_train, y_train):
    linreg = LinearRegression()
    linreg.fit(X_train, y_train)

    print("\n回归模型参数：")
    print(f"截距 θ0: {linreg.intercept_[0]:.4f}")
    coeff_df = pd.DataFrame(
        linreg.coef_.flatten(),
        index=X_train.columns,
        columns=['系数']
    )
    print(coeff_df)

    return linreg


def EvaluationModel(linreg, X_test, y_test):
    y_pred = linreg.predict(X_test)

    mse = mean_squared_error(y_test, y_pred)
    rmse = np.sqrt(mse)
    print("\n模型评估指标：")
    print(f"均方误差 (MSE): {mse:.4f}")
    print(f"均方根误差 (RMSE): {rmse:.4f}")

    return y_pred



def Visualization(y_test, y_pred):
    plt.figure(figsize=(8, 6))
    plt.scatter(y_test, y_pred, alpha=0.5)
    plt.plot([y_test.min(), y_test.max()],
             [y_test.min(), y_test.max()],
             'k--', lw=2, label='理想拟合线')
    plt.xlabel("实际值 (Measured)", fontsize=12)
    plt.ylabel("预测值 (Predicted)", fontsize=12)
    plt.title("实际值 vs 预测值", fontsize=14)
    plt.legend()
    plt.grid(True)
    plt.show()


if __name__ == "__main__":
    # 数据路径（需修改为实际路径）
    data_path = r"C:\pythonProject1\Folds5x2_pp.csv"

    # 数据预处理
    X_train, X_test, y_train, y_test = getTrainSetAndTestSet(data_path)
    # 模型训练
    model = TrainLinearRegression(X_train, y_train)

    # 模型评估
    y_pred = EvaluationModel(model, X_test, y_test)

    # 结果可视化
    Visualization(y_test, y_pred)